#!/usr/bin/env python
# encoding: utf-8
"""
wordstore.py

Copyright (c) 2008 William A. Carrel. All rights reserved.
"""

from google.appengine.ext import db

import logging
import pickle

class WordStoreShard(db.Model):
  trigram = db.StringProperty(required=True)
  pickle = db.BlobProperty(required=True)

class BadTrigramList(db.Model):
  bad_trigrams = db.StringListProperty()

ASCII_SMALL_A = 97

dirty_trigrams = {}

trigram_cache = {}

empty_trigrams = {}

def IdxFromChar(char):
  return ord(char)-ASCII_SMALL_A

def StringFromTrigramArray(array):
  return "%s%s%s" % (chr(array[0]+ASCII_SMALL_A),
                     chr(array[1]+ASCII_SMALL_A),
                     chr(array[2]+ASCII_SMALL_A),)

def FetchWordstore(trigram):
  """Fetch the wordstore for a given trigram."""
  # Get bad trigram list
  global empty_trigrams
  if not empty_trigrams:
    bad_trigram_list = BadTrigramList.get_or_insert("badtrigrams")
    for bad_trigram in bad_trigram_list.bad_trigrams:
      empty_trigrams[bad_trigram] = True

  if trigram in empty_trigrams:
    return None

  if trigram in trigram_cache:
    return trigram_cache[trigram]

  wordstore = {}
  # Fetch the wordstore shard
  # TODO(wac): Fix this so we use keys instead.
  raw_shard = WordStoreShard.get_by_key_name(trigram)
  # Try to regenerate raw_shard.
  if not raw_shard:
    query = WordStoreShard.all()
    query.filter("trigram =", trigram)
    old_shards = query.fetch(1)
    if old_shards and old_shards[0]:
      logging.info("Converting shard %s", trigram)
      old_shard = old_shards[0]
      raw_shard = WordStoreShard(trigram=old_shard.trigram,
                                 pickle=old_shard.pickle,
                                 key_name=trigram)
      db.delete(old_shard)
      raw_shard.put()
  # Nothing there, time to give up.
  if not raw_shard:
    empty_trigrams[trigram] = True
    return None
  wordstore.update(pickle.loads(raw_shard.pickle))
  trigram_cache[trigram] = wordstore
  if not wordstore:
    # This is a lame trigram, put it in the bad list if it isn't already
    empty_trigrams[trigram] = True
  return wordstore

class WordStoreWalk(object):
  def __init__(self, wordstore=None, with_create=False, trigram=[]):
    self.mark = 0
    self.wordstore = wordstore
    self.with_create = with_create
    self.trigram = trigram
    if not self.wordstore and len(self.trigram) == 3:
      # Need to pickup the dirty trigram if there is one waiting
      global dirty_trigrams
      if StringFromTrigramArray(trigram) in dirty_trigrams:
        self.wordstore = dirty_trigrams[StringFromTrigramArray(trigram)]
      else:
        self.wordstore = FetchWordstore(StringFromTrigramArray(trigram))
      if not self.wordstore:
        self.wordstore = {}
      if self.with_create:
        dirty_trigrams[StringFromTrigramArray(trigram)] = self.wordstore

  def AddWordHere(self, word):
    if self.wordstore is not None:
      self.wordstore['word'] = word
    else:
      logging.debug("Not adding word '%s', it's too short.", word)

  def Descend(self, word):
    i = IdxFromChar(word[0])
    return self.DescendInt(i)

  def DescendInt(self, i):
    # We're not far enough in to have loaded our wordstore shard yet
    if self.wordstore is None:
      new_trigram = self.trigram[:]
      new_trigram.append(i)
      return WordStoreWalk(with_create=self.with_create, trigram=new_trigram)
    if i not in self.wordstore:
      if not self.with_create:
        return None
      self.wordstore[i] = {}
    return WordStoreWalk(self.wordstore[i], self.with_create)

def DumpPrefix(prefix):
  walk = WordStoreWalk()
  
  for letter in prefix:
    walk = walk.Descend(letter)
    if not walk:
      return None
  
  return walk.wordstore

def AddWord(word):
  walk = WordStoreWalk(with_create=True)
  for letter in word:
    walk = walk.Descend(letter)
  walk.AddWordHere(word)

def WriteBadTrigrams():
  bad_trigram_list = BadTrigramList.get_or_insert("badtrigrams")
  
  global empty_trigrams
  for trigram in bad_trigram_list.bad_trigrams:
    empty_trigrams[trigram] = True
  
  bad_trigram_list.bad_trigrams = [i for i in empty_trigrams if empty_trigrams[i]]
  bad_trigram_list.put()


def FlushChanges():
  global dirty_trigrams  # Load in the data we're going to be using

  for trigram in dirty_trigrams:
    query = WordStoreShard.all()
    query.filter('trigram =', trigram)
    results = query.fetch(1)
    if results:
      shard = results[0]
    else:
      shard = WordStoreShard(trigram=trigram, pickle='uninitialized')
      shard.trigram = trigram
      if trigram in empty_trigrams:
        del empty_trigrams[trigram]
    shard.pickle = pickle.dumps(dirty_trigrams[trigram])
    shard.put()
    trigram_cache[trigram] = dirty_trigrams[trigram]
  dirty_trigrams = {}
  # TODO(wac): Write out the empty trigram list.
  